import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
import os
import numpy as np

# --- 设置 ---
# 确保图片输出目录存在
output_dir = 'Pictures'
os.makedirs(output_dir, exist_ok=True)

# 加载特征数据
try:
    df = pd.read_csv('features.csv')
except FileNotFoundError:
    print("错误: features.csv 文件未找到。请先运行特征提取脚本。")
    exit()

# --- 预处理：为统一坐标轴做准备 ---
# 找到所有包中的最大命中数，以统一X轴
max_hits = df['hit_rules_count'].max()
x_ticks = range(max_hits + 1)

# --- 1. 恶意包分析 ---
print("="*50)
print("恶意包规则命中分析 (Label=1)")
print("="*50)

malicious_df = df[df['label'] == 1].copy()

if malicious_df.empty:
    print("数据集中没有恶意包。")
else:
    total_malicious = len(malicious_df)
    packages_hit_malicious = malicious_df[malicious_df['hit_rules_count'] > 0]
    coverage_rate = len(packages_hit_malicious) / total_malicious if total_malicious > 0 else 0

    print(f"恶意包总数: {total_malicious}")
    print(f"被规则命中的恶意包数量: {len(packages_hit_malicious)}")
    print(f"规则集在恶意包上的覆盖率: {coverage_rate:.2%}")

    # 准备绘图数据，确保所有x值都有对应的y值（即使是0）
    hit_counts_malicious = malicious_df['hit_rules_count'].value_counts().reindex(x_ticks, fill_value=0)
    print("\n规则命中数分布详情:")
    print(hit_counts_malicious[hit_counts_malicious > 0]) # 只打印有命中的部分

    # --- 可视化：恶意包命中分布 ---
    plt.figure(figsize=(12, 7))
    ax_malicious = sns.barplot(x=hit_counts_malicious.index, y=hit_counts_malicious.values, palette="Reds_r")
    
    # 使用对数坐标轴
    ax_malicious.set_yscale("log")
    
    plt.title('Distribution of Hit Rules Count for Malicious Packages (Log Scale)', fontsize=16)
    plt.xlabel('Number of Rules Hit', fontsize=12)
    plt.ylabel('Number of Malicious Packages (Log Scale)', fontsize=12)
    
    # 调整X轴刻度，避免过于密集
    if max_hits > 20:
        plt.xticks(ticks=[i for i in x_ticks if i % 5 == 0 or i == 1], rotation=45)
    else:
        plt.xticks(ticks=x_ticks, rotation=45)

    # 在柱子上显示数值
    for p in ax_malicious.patches:
        height = p.get_height()
        if height > 0:
            ax_malicious.annotate(f'{int(height)}', (p.get_x() + p.get_width() / 2., height),
                                  ha='center', va='bottom', xytext=(0, 5), textcoords='offset points')

    plt.tight_layout()
    malicious_plot_path = os.path.join(output_dir, 'malicious_rule_coverage.png')
    plt.savefig(malicious_plot_path, dpi=300)
    print(f"\n恶意包命中分布图已保存至: {malicious_plot_path}")


# --- 2. 良性包分析 ---
print("\n" + "="*50)
print("良性包规则命中分析 (Label=0)")
print("="*50)

benign_df = df[df['label'] == 0].copy()

if benign_df.empty:
    print("数据集中没有良性包。")
else:
    total_benign = len(benign_df)
    packages_hit_benign = benign_df[benign_df['hit_rules_count'] > 0]
    fp_rate_by_rules = len(packages_hit_benign) / total_benign if total_benign > 0 else 0

    print(f"良性包总数: {total_benign}")
    print(f"被规则命中的良性包数量 (潜在误报源): {len(packages_hit_benign)}")
    print(f"规则集在良性包上的潜在误报率: {fp_rate_by_rules:.4%}")

    # 准备绘图数据，确保所有x值都有对应的y值（即使是0）
    hit_counts_benign = benign_df['hit_rules_count'].value_counts().reindex(x_ticks, fill_value=0)
    print("\n规则命中数分布详情:")
    print(hit_counts_benign[hit_counts_benign > 0]) # 只打印有命中的部分

    # --- 可视化：良性包命中分布 ---
    plt.figure(figsize=(12, 7))
    ax_benign = sns.barplot(x=hit_counts_benign.index, y=hit_counts_benign.values, palette="Greens_r")

    # 使用对数坐标轴
    ax_benign.set_yscale("log")

    plt.title('Distribution of Hit Rules Count for Benign Packages (Log Scale)', fontsize=16)
    plt.xlabel('Number of Rules Hit', fontsize=12)
    plt.ylabel('Number of Benign Packages (Log Scale)', fontsize=12)

    # 调整X轴刻度，避免过于密集
    if max_hits > 20:
        plt.xticks(ticks=[i for i in x_ticks if i % 5 == 0 or i == 1], rotation=45)
    else:
        plt.xticks(ticks=x_ticks, rotation=45)

    # 在柱子上显示数值
    for p in ax_benign.patches:
        height = p.get_height()
        if height > 0:
            ax_benign.annotate(f'{int(height)}', (p.get_x() + p.get_width() / 2., height),
                                  ha='center', va='bottom', xytext=(0, 5), textcoords='offset points')

    plt.tight_layout()
    benign_plot_path = os.path.join(output_dir, 'benign_rule_coverage.png')
    plt.savefig(benign_plot_path, dpi=300)
    print(f"\n良性包命中分布图已保存至: {benign_plot_path}")

print("\n分析完成。")